
"""
Contrastive Learning for Space Alignment (Paper Section 2.4.3)
Implements Equation 13: Cross-modal contrastive loss
"""

import torch
import torch.nn as nn

class CrossModalContrastiveLoss(nn.Module):
    def __init__(self, temperature=0.1):
        """
        Args:
            temperature: τ parameter in Equation 13
        """
        super().__init__()
        self.temp = temperature
        self.cross_entropy = nn.CrossEntropyLoss()

    def forward(self, behavior_emb, text_emb):
        """
        Args:
            behavior_emb: Behavior embeddings [B, D]
            text_emb: Text embeddings [B, D]
        Returns:
            contrastive_loss: NCE-style loss value
        """
        # Normalization
        behavior_norm = F.normalize(behavior_emb, p=2, dim=-1)
        text_norm = F.normalize(text_emb, p=2, dim=-1)
        
        # Similarity matrix
        sim_matrix = torch.mm(behavior_norm, text_norm.T) / self.temp
        
        # Positive pairs are diagonal entries
        batch_size = behavior_emb.size(0)
        labels = torch.arange(batch_size).to(behavior_emb.device)
        
        # Symmetric loss calculation
        loss_bt = self.cross_entropy(sim_matrix, labels)  # Behavior vs Text
        loss_tb = self.cross_entropy(sim_matrix.T, labels) # Text vs Behavior
        return (loss_bt + loss_tb) / 2
